package rbtree.non_recursion_3;

import rbtree.non_recursion.Color;

/**
 * @author xing
 * @date 2023/1/12 10:51:53
 */
public class RBTree<K extends Comparable<K>, V> {

    public final RBNode<K, V> nilNode = new RBNode<>();

    private RBNode<K, V> root;

    public RBTree() {
    }

    public RBTree(RBNode<K, V> root) {
        this.root = root;
    }

    public RBNode<K, V> getRoot() {
        return root;
    }

    public void insert(K key, V value) {
        if (key == null) {
            return;
        }
        RBNode<K, V> newNode = new RBNode<>(key, value);
        if (this.root == null) {
            this.root = newNode;
            this.root.color = RBColor.BLACK;
            return;
        }
        RBNode<K, V> temp = this.root;
        RBNode<K, V> newNodeParent = temp;
        while (temp != null) {
            newNodeParent = temp;
            if (temp.key.compareTo(newNode.key) < 0) {
                temp = temp.right;
            } else {
                temp = temp.left;
            }
        }
        newNodeParent.setLeftOrRightAndAsParent(newNode);
        insertFixUp(newNode);
    }

    private void insertFixUp(RBNode<K, V> node) {
        while (node != this.root) {
            RBNode<K, V> parent = node.parent;
            // 父为黑
            // zhu: 如果父节点是根节点，（没有祖父节点），根节点为黑色，直接 return , 所以下边一定有 grandParent
            if (isBlack(parent)) {
                return;
            }
            // 父为红
            RBNode<K, V> grandParent = parent.parent;
            // 父在右
            if (grandParent.key.compareTo(parent.key) < 0) {
                // 叔叔为红，变色
                if (isRed(grandParent.left)) {
                    changeColorByInsertFour(grandParent);
                    node = grandParent;
                }
                // 叔叔为黑 或者 为 Null (nil 节点为 黑)，旋转
                else {
                    // 在父右
                    if (parent.key.compareTo(node.key) < 0) {
                        // RR
                        leftRotate(grandParent);
                        changeColorByInsertThree(parent);
                        return;
                    }
                    // 在父左
                    else {
                        // RL
                        // 只需要对 parent 右旋，然后交给下一次循环
                        rightRotate(parent);
                    }
                }
            }
            // 父在左
            else {
                // 叔叔为红，变色
                if (isRed(grandParent.right)) {
                    changeColorByInsertFour(grandParent);
                    node = grandParent;
                }
                // 叔叔为黑 或者 为 Null (nil 节点为 黑)，旋转
                else {
                    // 在父右
                    if (parent.key.compareTo(node.key) < 0) {
                        // LR
                        // 只需要对 parent 左旋，然后交给下一次循环
                        leftRotate(parent);
                    }
                    // 在父左
                    else {
                        // LL
                        rightRotate(grandParent);
                        changeColorByInsertThree(parent);
                        return;
                    }
                }
            }
        }
        this.root.color = RBColor.BLACK;
    }

    public void delete(K key, V value) {
        // deletedNode 为最终要被删除的结点
        RBNode<K, V> deletedNode = findNode(key);
        if (deletedNode == null) {
            return;
        }
        // 该判定为寻找最终要删除的 Node . 
        // 左右子树均不为空， (zhi you 一个子结点时，不需要计算前驱结点，直接上移占他位置, 此时 deleteNode 不变 ) . 。
        if (deletedNode.left != null && deletedNode.right != null) {
            // 获取前驱结点
            RBNode<K, V> temp = getPrecursor(deletedNode);
            // 复制数据给原被删结点（父子关系，颜色等都不需要再改，方便）
            deletedNode.key = temp.key;
            deletedNode.value = temp.value;
            // 被删除结点转变为该前驱结点
            deletedNode = temp;
        }
        // 孩子结点， (用来作为 deleteFixup() 方法的参数，。标记删除节点的位置，。所以当没有子节点，使用 nil 结点替代，。)
        RBNode<K, V> childNode;
        // 获取孩子结点
        if (deletedNode.left != null) {
            childNode = deletedNode.left;
        } else if (deletedNode.right != null) {
            childNode = deletedNode.right;
        } else {
            // 左右都不存在，赋 nil Node
            childNode = nilNode;
        }
        // 删除 deletedNode，将子结点给其父结点
        childNode.parent = deletedNode.parent;
        // 父结点为 null，deletedNode 为根结点
        if (deletedNode.parent == null) {
            // 重设根结点 为孩子结点
            // 删除可能的 nilNode
            if (childNode == nilNode) {
                this.root = null;
                return;
            }
            this.root = childNode;
            // childNode 可能为红色，修改一下根结点颜色
            this.root.color = RBColor.BLACK;
            // 结束，返回
            return;
        } 
        // 父结点不为 null，被删结点不是根结点
        else {
            // 将 deletedNode 替换为 childNode
            deletedNode.parent.replaceChild(deletedNode, childNode);
        }
        // 被删结点为黑结点
        if (isBlack(deletedNode)) {
            // 做术后修复
            deleteFixup(childNode);
        } else {
            // 不是黑节点，删除可能的 nilNode
            deleteIfNilNode(childNode);
        }
        // 看似有些代码前面设置，后边删除，很多次一举，但正是如此，代码才这么少，
        // 不然中间有可能有情况就判断, 代码量会多很多
    }

    private void deleteFixup(RBNode<K, V> node) {
        while (node != this.root && node.color == RBColor.BLACK) {
            RBNode<K, V> parent = node.parent;
            if (node == parent.left) {
                // 兄弟红
                if (isRed(parent.right)) {
                    // 三结点变换
                    parent.color = RBColor.RED;
                    parent.right.color = RBColor.BLACK;
                    leftRotate(parent);
                }
                // 兄弟黑
                else {
                    // 兄弟右子红
                    if (isRed(parent.right.right)) {
                        parent.right.color = parent.color;
                        parent.color = RBColor.BLACK;
                        leftRotate(parent);
                        deleteIfNilNode(node);
                        // node = parent.parent;
                        // 平衡完毕，提前结束循环。
                        return;
                    }
                    // 兄弟左子红
                    else if (isRed(parent.right.left)){
                        parent.right.color = RBColor.RED;
                        parent.right.left.color = RBColor.BLACK;
                        rightRotate(parent);
                    }
                    // 兄弟子均黑
                    else {
                        parent.right.color = RBColor.RED;
                        deleteIfNilNode(node);
                        node = parent;
                    }
                }
            } else {
                // 兄弟红
                if (isRed(parent.left)) {
                    // 三结点变换
                    parent.color = RBColor.RED;
                    parent.left.color = RBColor.BLACK;
                    rightRotate(parent);
                }
                // 兄弟黑
                else {
                    // 兄弟左子红
                    if (isRed(parent.left.left)) {
                        parent.left.color = parent.color;
                        parent.color = RBColor.BLACK;
                        rightRotate(parent);
                        deleteIfNilNode(node);
                        // node = parent.parent;
                        // 平衡完毕，提前结束循环。
                        return;
                    }
                    // 兄弟右子红
                    else if (isRed(parent.left.right)){
                        parent.left.color = RBColor.RED;
                        parent.left.right.color = RBColor.BLACK;
                        leftRotate(parent);
                    }
                    // 兄弟子均黑
                    else {
                        parent.left.color = RBColor.RED;
                        deleteIfNilNode(node);
                        node = parent;
                    }
                }
            }
        }
        node.color = RBColor.BLACK;
    }

    private void deleteIfNilNode(RBNode<K, V> node) {
        if (node == nilNode) {
            if (node == node.parent.left) {
                node.parent.left = null;
            } else {
                node.parent.right = null;
            }
            nilNode.parent = null;
        }
    }

    // // delete() 方法中做了 left != null 的判断，所以这里不使用 这种方式 计算 precursor    . 。
    // public RBNode<K, V> getPrecursor(RBNode<K, V> node) {
    //     RBNode<K, V> precursor;
    //     if ((precursor = node.left) == null) {
    //         return node;
    //     }
    //     while (precursor.right != null) {
    //         precursor = precursor.right;
    //     }
    //     return precursor;
    // }

    private RBNode<K, V> getPrecursor(RBNode<K, V> node) {
        RBNode<K, V> left = node.left;
        while (left.getRight() != null) {
            left = left.right;
        }
        return left;
    }

    private RBNode<K, V> findNode(K key) {
        if (key == null) {
            return null;
        }
        RBNode<K, V> temp = this.root;
        int cmp;
        while (temp != null) {
            cmp = temp.key.compareTo(key);
            if (cmp < 0) {
                temp = temp.right;
            } else if (cmp > 0) {
                temp = temp.left;
            } else {
                return temp;
            }
        }
        return null;
    }

    private void changeColorByInsertFour(RBNode<K, V> node) {
        node.color = RBColor.RED;
        node.left.color = RBColor.BLACK;
        node.right.color = RBColor.BLACK;
    }

    private void changeColorByInsertThree(RBNode<K, V> node) {
        node.color = RBColor.BLACK;
        node.left.color = RBColor.RED;
        node.right.color = RBColor.RED;
    }

    private boolean isBlack(RBNode<K, V> node) {
        // nil 节点是特殊的 final 节点，不为 Null ，但是 为 黑色。
        // 理论上为 Null 依然应该认为是 黑色 ，。（并且 在删除逻辑中，会用到。）
        return node == null || node.color == RBColor.BLACK;
    }

    private boolean isRed(RBNode<K, V> node) {
        // return !isBlack(node);
        return node != null && node.color == RBColor.RED ;
    }


    private void leftRotate(RBNode<K, V> root) {
        if (root.right == null) {
            return;
        }
        RBNode<K, V> right = root.right;
        if (root.parent != null) {
            root.parent.setLeftOrRightAndAsParent(right);
        } else {
            right.parent = null;
            this.root = right;
        }
        root.setRightAndAsParent(right.left);
        right.setLeftAndAsParent(root);
    }

    private void rightRotate(RBNode<K, V> root) {
        if (root.left == null) {
            return;
        }
        RBNode<K, V> left = root.left;
        if (root.parent != null) {
            root.parent.setLeftOrRightAndAsParent(left);
        } else {
            left.parent = null;
            this.root = left;
        }
        root.setLeftAndAsParent(left.right);
        left.setRightAndAsParent(root);
    }

    /**
     * 红黑树结点颜色
     */
    protected enum RBColor {
        RED,
        BLACK
    }

    /**
     * 红黑树结点
     */
    public static class RBNode<K extends Comparable<K>, V> {
        private RBNode<K, V> parent;
        private RBNode<K, V> left;
        private RBNode<K, V> right;
        private RBColor color;
        private K key;
        private V value;

        private RBNode() {
            this.color = RBColor.BLACK;
        }

        public RBNode(K key, V value) {
            this.color = RBColor.RED;
            this.key = key;
            this.value = value;
        }

        public void setLeftAndAsParent(RBNode<K, V> left) {
            this.left = left;
            if (left != null) {
                left.parent = this;
            }
        }

        public void setRightAndAsParent(RBNode<K, V> right) {
            this.right = right;
            if (right != null) {
                right.parent = this;
            }
        }

        public void setLeftOrRightAndAsParent(RBNode<K, V> node) {
            if (node == null) {
                throw new IllegalArgumentException("node 不能为 null，无法判断插入位置");
            }
            if (this.key.compareTo(node.key) < 0) {
                setRightAndAsParent(node);
            } else {
                setLeftAndAsParent(node);
            }
        }

        public void replaceChild(RBNode<K, V> child, RBNode<K, V> replacement) {
            if (child == this.left) {
                setLeftAndAsParent(replacement);
            } else if (child == this.right) {
                setRightAndAsParent(replacement);
            } else {
                throw new IllegalArgumentException("参数有误，没有此子结点，无法替换");
            }
        }

        public RBNode<K, V> getLeft() {
            return left;
        }

        public RBNode<K, V> getRight() {
            return right;
        }

        public RBColor getColor() {
            return color;
        }

        public K getKey() {
            return key;
        }

    }
}
